#!/usr/bin/python

import ast
import random
import sys

# Will turn every integer constant into a sum of a bunch of random integers
class IntInflationObfuscation(ast.NodeTransformer):
    def __init__(self, min_length, max_length, min_value, max_value):
        self.min_length = min_length
        self.max_length = max_length
        self.min_value = min_value
        self.max_value = max_value
    
    def get_const(self, value: int):
        if value >= 0:
            return f"-(-({value}))"

        return f"-({abs(value)})"

    def visit_Constant(self, node: ast.Constant):
        if not (type(node.value) is int): return node
        
        length = random.randint(self.min_length, self.max_length)
        values = [random.randint(self.min_value, self.max_value) for _ in range(length-1)]

        expr_string = "+".join([self.get_const(v) for v in values]) + f"+{self.get_const(node.value - sum(values))}"
                
        return ast.parse(expr_string, mode="eval").body


# EXPERIMENTAL: replaces all identifiers with gibberish
# WARNING: this does not keep track of identifier `path`
# for example: if a user defined class `A` has a field of `foo` and an external class `B` has a field `foo` then any access to `B.foo` will be treated as if it is an access to `A.foo` and be obfuscated, which will break the program
# Same goes for variable names, function names and any other identifiers
# CURRENTLY THE UNSAFE BEHAVIOR IS COMMENTED OUT

class VariableNameObfuscation(ast.NodeTransformer):
    unicode_start = 0x41 #0x0388
    unicode_end = 0x5a #0x03ff 

    class identifier_analyzer(ast.NodeVisitor):
        def __init__(self):
            self.names: set[str] = set()
                   
        def visit_For(self, node: ast.For):
            if type(node.target) is ast.Name:
                self.names.add(node.target.id)

            return self.generic_visit(node)

        def visit_FunctionDef(self, node: ast.FunctionDef):
            self.names.add(node.name)

            for arg in node.args.posonlyargs:
                self.names.add(arg.arg)

            for arg in node.args.kwonlyargs:
                self.names.add(arg.arg)
            
            for arg in node.args.args:
                self.names.add(arg.arg)

            if node.args.vararg != None:
                self.names.add(node.args.vararg.arg)

            if node.args.kwarg != None:
                self.names.add(node.args.kwarg.arg)

            self.generic_visit(node)
        
        def visit_ClassDef(self, node: ast.ClassDef):
            self.names.add(node.name)

            self.generic_visit(node)
        
        def visit_Assign(self, node: ast.Assign):
            for target in node.targets:
                if type(target) is ast.Name:
                    self.names.add(target.id)
            
            return self.generic_visit(node)
        
        def visit_Import(self, node: ast.Import):
            for n in node.names:
                if n.asname != None:
                    self.names.add(n.asname)
                else:
                    self.names.add(n.name)
        
        def visit_ImportFrom(self, node: ast.ImportFrom):
            for n in node.names:
                self.names.add(n.name)
        
        def get_names(self):
            return self.names


    def __init__(self, min_len: int, max_len: int, tree: ast.AST):
        ia = self.identifier_analyzer() 
        ia.visit(tree)

        self.names: set[str] = ia.get_names() 
        self.aliases: dict[str, str] = {}
        self.min_len = min_len
        self.max_len = max_len
        self.inside_class = False
    
    def alias(self, name: str) -> str:
        if not (name in self.names): return name

        # special case
        if name.startswith("__"): return name

        if not (name in self.aliases.keys()):
            length = random.randint(self.min_len, self.max_len)
            new_name = "".join([chr(random.randint(self.unicode_start, self.unicode_end-1)) for i in range(length)])
            self.aliases[name] = new_name

        return self.aliases[name]

    def visit_Name(self, node: ast.Name):
        return ast.Name(id=self.alias(node.id), ctx=node.ctx)

    # if this is commented out, it should be safe
    ''' 
    def visit_Attribute(self, node: ast.Attribute):
        self.generic_visit(node)
        result = ast.Attribute(value=node.value, attr=self.alias(node.attr), ctx=node.ctx)
        
        return result
    '''
        
    def visit_FunctionDef(self, node: ast.FunctionDef):
        self.generic_visit(node)
        
        new_args = ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], vararg=None, kwarg=None, kw_defaults=node.args.kw_defaults, defaults=node.args.defaults)
        for poa in node.args.posonlyargs:
            new_args.posonlyargs.append(ast.arg(arg=self.alias(poa.arg), annotation=poa.annotation, type_comment=poa.type_comment))
        
        for a in node.args.args:
            new_args.args.append(ast.arg(arg=self.alias(a.arg), annotation=a.annotation, type_comment=a.type_comment))

        for koa in node.args.posonlyargs:
            new_args.kwonlyargs.append(ast.arg(arg=self.alias(koa.arg), annotation=koa.annotation, type_comment=koa.type_comment))
       
        if node.args.vararg != None:
            new_args.vararg = ast.arg(arg=self.alias(node.args.vararg.arg), annotation=node.args.vararg.annotation, type_comment=node.args.vararg.type_comment)

        if node.args.kwarg != None:
            new_args.kwarg = ast.arg(arg=self.alias(node.args.kwarg.arg), annotation=node.args.kwarg.annotation, type_comment=node.args.kwarg.type_comment)

        if self.inside_class: 
            result = ast.FunctionDef(name=node.name, args=new_args, body=node.body, decorator_list=node.decorator_list, returns=node.returns, type_params=node.type_params)
        else:
            result = ast.FunctionDef(name=self.alias(node.name), args=new_args, body=node.body, decorator_list=node.decorator_list, returns=node.returns, type_params=node.type_params)
        
        return result

    def visit_ClassDef(self, node: ast.ClassDef):
        # nested
        if self.inside_class:
            self.generic_visit(node)
        # not nested
        else:
            self.inside_class = True
            self.generic_visit(node)
            self.inside_class = False
        
        if not self.inside_class:
            result = ast.ClassDef(name=self.alias(node.name), body=node.body, bases=node.bases, keywords=node.keywords, decorator_list=node.decorator_list, type_params=node.type_params)
        else:
            # don't obfuscate nested classes, or it breaks completely
            result = node
            
        return result

    def visit_Import(self, node: ast.Import):
        result = ast.Import(names=[])
        
        for name in node.names:
            if name.asname != None:
                result.names.append(ast.alias(name=name.name, asname=self.alias(name.asname)))
            else:
                result.names.append(ast.alias(name=name.name, asname=self.alias(name.name)))
        
        return result

    def visit_ImportFrom(self, node: ast.ImportFrom):
        result = ast.ImportFrom(module=node.module, level=node.level, names=[])
        
        for name in node.names:
            result.names.append(ast.alias(name=name.name, asname=self.alias(name.name)))
        
        return result

# This replaces attribute reads with `getattr` calls
# and attribute sets with `setattr` calls
class AttributeAccessObfuscation(ast.NodeTransformer):
    def visit_Attribute(self, node: ast.Attribute):
        self.generic_visit(node)
        if type(node.ctx) is ast.Store: return node

        obj = node.value
        attr = ast.Constant(value=node.attr)
        
        return ast.Call(func=ast.Name(id="getattr", ctx=ast.Load()), args=[obj, attr], keywords=[])
    
    def visit_Assign(self, node: ast.Assign):
        self.generic_visit(node)

        result = []
        for target in node.targets:
            if type(target) is ast.Attribute:
                result.append(
                    ast.Expr(value=
                        ast.Call(func=ast.Name(id="setattr", ctx=ast.Load()), args=[target.value, ast.Constant(target.attr), node.value], keywords=[])
                    )
                )
            else:
                result.append(
                    ast.Assign(targets=[target], value=node.value)
                )
        
        return result

# This will turn every string constant into a sum of `chr()`
class StringInflationObfuscation(ast.NodeTransformer):
    def get_call_for_char(self, char):
        return f"chr({ord(char)})" 

    # No support for f-strings
    def visit_JoinedStr(self, node: ast.JoinedStr):
        return node

    def visit_Constant(self, node: ast.Constant):
        if not (type(node.value) is str): return node

        expr_string = "+".join([self.get_call_for_char(c) for c in node.value]) 
                
        return ast.parse(expr_string, mode="eval").body

# This will just hide all imports behind `exec`s
class HideImportsObfuscation(ast.NodeTransformer):
    def visit_Import(self, node: ast.Import):
        exec_string = f"exec('{ast.unparse(node)}')"
        return ast.parse(exec_string).body
    
    def visit_ImportFrom(self, node: ast.ImportFrom):
        exec_string = f"exec('{ast.unparse(node)}')"
        return ast.parse(exec_string).body
    

# This will wrap any built-in usage in a `builtins.` 
# Intended to be used together with AttributeAccessObfuscation
# This crashes and burns if user code has redefined built-in names
# For example: a `def input()`
# PLEASE DON'T DO THAT THAT IS DUMB
class BuiltinsObfuscation(ast.NodeTransformer):
    builtins_list = [
        'ArithmeticError', 
        'AssertionError', 
        'AttributeError', 
        'BaseException', 
        'BaseExceptionGroup', 
        'BlockingIOError', 
        'BrokenPipeError', 
        'BufferError', 
        'BytesWarning', 
        'ChildProcessError', 
        'ConnectionAbortedError', 
        'ConnectionError', 
        'ConnectionRefusedError', 
        'ConnectionResetError', 
        'DeprecationWarning', 
        'EOFError', 
        'Ellipsis', 
        'EncodingWarning', 
        'EnvironmentError', 
        'Exception', 
        'ExceptionGroup', 
        'False', # builtins.False doesn't work, but getattr(builtins, "False") does lmao
        'FileExistsError', 
        'FileNotFoundError', 
        'FloatingPointError', 
        'FutureWarning', 
        'GeneratorExit', 
        'IOError', 
        'ImportError', 
        'ImportWarning', 
        'IndentationError', 
        'IndexError', 
        'InterruptedError', 
        'IsADirectoryError', 
        'KeyError', 
        'KeyboardInterrupt', 
        'LookupError', 
        'MemoryError', 
        'ModuleNotFoundError', 
        'NameError', 
        'None', # Doesn't work as well :///
        'NotADirectoryError', 
        'NotImplemented', 
        'NotImplementedError', 
        'OSError', 
        'OverflowError', 
        'PendingDeprecationWarning', 
        'PermissionError', 
        'ProcessLookupError', 
        'RecursionError', 
        'ReferenceError', 
        'ResourceWarning', 
        'RuntimeError', 
        'RuntimeWarning', 
        'StopAsyncIteration', 
        'StopIteration', 
        'SyntaxError', 
        'SyntaxWarning', 
        'SystemError', 
        'SystemExit', 
        'TabError', 
        'TimeoutError', 
        'True', # no workey
        'TypeError', 
        'UnboundLocalError', 
        'UnicodeDecodeError', 
        'UnicodeEncodeError', 
        'UnicodeError', 
        'UnicodeTranslateError', 
        'UnicodeWarning', 
        'UserWarning', 
        'ValueError', 
        'Warning', 
        'ZeroDivisionError', 
        # Commenting these out just in case '__build_class__', '__debug__', '__doc__', '__import__', '__loader__', '__name__', '__package__', '__spec__', 
        'abs', 
        'aiter', 
        'all', 
        'anext', 
        'any', 
        'ascii', 
        'bin', 
        'bool', 
        'breakpoint', 
        'bytearray', 
        'bytes', 
        'callable', 
        'chr', 
        'classmethod', 
        'compile', 
        'complex', 
        'copyright', 
        'credits', 
        'delattr', 
        'dict', 
        'dir', 
        'divmod', 
        'enumerate', 
        'eval', 
        'exec', 
        'exit', 
        'filter', 
        'float', 
        'format', 
        'frozenset', 
        'getattr', 
        'globals', 
        'hasattr', 
        'hash', 
        'help', 
        'hex', 
        'id', 
        'input', 
        'int', 
        'isinstance', 
        'issubclass', 
        'iter', 
        'len', 
        'license', 
        'list', 
        'locals', 
        'map', 
        'max', 
        'memoryview', 
        'min', 
        'next', 
        'object',
        'oct', 
        'open', 
        'ord', 
        'pow', 
        'print', 
        'property', 
        'quit', 
        'range', 
        'repr', 
        'reversed', 
        'round', 
        'set', 
        'setattr', 
        'slice', 
        'sorted', 
        'staticmethod', 
        'str', 
        'sum', 
        'super', 
        'tuple', 
        'type', 
        'vars', 
        'zip'
    ]
    
    def __init__(self):
        self.builtins_imported = False
    
    def visit_Import(self, node: ast.Import):
        for name in node.names:
            if name.name == "builtins":
                self.builtins_imported = True
                return node
        
        return node

    # first, add a `builtins` import
    def visit_Module(self, node: ast.Module):
        self.generic_visit(node)
        if self.builtins_imported: return node
        
        imp_expr = ast.parse("import builtins").body[0]
        
        return ast.Module(body=[imp_expr]+node.body, type_ignores=node.type_ignores)


    # replacing all builtins with attribute accesses
    def visit_Name(self, node: ast.Name):
        if not (node.id in self.builtins_list): return node
        
        attr_expr = ast.parse(f"builtins.{node.id}", mode="eval").body
        
        return attr_expr

    # for False, True and None
    def visit_Constant(self, node: ast.Constant):
        if not (type(node.value) in [type(True), type(None)]): return node
        
        getattr_expr = ast.parse(f"builtins.getattr(builtins, '{node.value}')", mode="eval").body
        
        return getattr_expr


if __name__ == "__main__":
    obfuscation_table = {
        "i": IntInflationObfuscation,
        "s": StringInflationObfuscation,
        "v": VariableNameObfuscation,
        "a": AttributeAccessObfuscation,
        "m": HideImportsObfuscation,
        "b": BuiltinsObfuscation,
    }    
    
    default_values = {
        "i": [5, 10, -1000, 1000],
        "v": [3, 25],
    }
    
    tree = None
    with open(sys.argv[2]) as f:
        contents = f.read()
        tree = ast.parse(contents)
    
    program = sys.argv[1]
    for p in program:
        obf = obfuscation_table[p]
        
        values: None | list = None
        if p in default_values:
            values = default_values[p].copy()
        
        if p == "v" and values != None:
            values.append(tree)
        
        if values != None:
            tree = ast.fix_missing_locations(obf(*values).visit(tree))
        else:
            tree = ast.fix_missing_locations(obf().visit(tree))
        
    print(ast.unparse(tree))    

